"""
构建法律案例的向量数据库，将法律案例文本转换为可用于相似性检索的向量索引。

1. 数据加载与合并
从两个 JSONL 文件（train_data.jsonl和rest_data.jsonl）中读取法律案例数据，每个案例包含：
input：案情描述文本（如犯罪事实、审理过程等）
output：判决结果（罪名、刑期、罚金、犯罪嫌疑人等信息）
将两个文件的案例数据合并为一个列表，作为构建数据库的原始数据。

2. 数据预处理与文档构造
过滤掉案情描述过长（超过 2000 字符）的案例，避免处理效率问题。
对每个案例的判决结果（output）进行格式化处理（通过predict.accusation_json2str转换为字符串，便于后续展示）。
为每个案例构建Document对象（LangChain 中的文档结构）：
page_content：存储案情描述文本（用于转换为向量，参与检索）
metadata：存储附加信息（判决结果result和罪名category，不参与向量计算，仅作为检索结果的辅助信息）
若一个案例涉及多个罪名（accusation_list），则为每个罪名创建一个独立的Document，确保后续可按罪名筛选检索。

3. 向量数据库构建
使用自定义的CustomEmbeddings（基于 BGE 模型）将所有案例的page_content（案情描述）转换为归一化的向量。
通过FAISS.from_documents将向量与文档关联，构建 FAISS 向量索引（一种高效的相似性搜索工具）。
将构建好的向量数据库（db）通过pickle序列化并保存到本地文件law_db，便于后续加载使用。
"""
from langchain_community.vectorstores import FAISS

import pickle
import json
from langchain.docstore.document import Document
from .Embedding import CustomEmbeddings
import predict

embeddings = CustomEmbeddings()
documents = []
with open("data/train_data.jsonl", encoding="utf-8") as f:
	lines1 = [json.loads(s.strip()) for s in f.readlines()]
with open("data/rest_data.jsonl", encoding="utf-8") as f:
	lines2 = [json.loads(s.strip()) for s in f.readlines()]
lines = lines1 + lines2
"""
{
 'input': '昌宁县人民检察院指控，2014 年 4 月 19 日下午 16 时许，被告人段某驾拖车经过鸡飞乡澡塘街子，时逢堵车，段某将车停在“冰凉一夏”冷饮店门口，被害人王某的侄子王2某示意段某靠边未果，后上前敲打车门让段某离开，段某遂驾车离开，但对此心生怨愤。同年4月21日22时许，被告人段某酒后与其妻子王1某一起准备回家，走到鸡飞乡澡塘街富达通讯手机店门口时停下，段某进入手机店内对被害人王某进行吼骂，紧接着从手机店出来拿得一个石头又冲进手机店内朝王某头部打去，致王某右额部粉碎性骨折、右眼眶骨骨折。经鉴定，被害人王某此次损伤程度为轻伤一级。',
 'output': {'是否死刑': False,
  '有期徒刑': 12,
  '是否无期': False,
  '罪名': ['故意伤害'],
  '罚金': 0,
  '犯罪嫌疑人': ['段某']}
  }
"""

# 构建法律案件的文档数据库
for data in lines:
	text = data["input"]
	if len(text) > 2000:
		continue
	result = predict.accusation_json2str(data["output"])
	
	# 获取对应案件的真实罪名列表
	accusation_list = data["output"]["罪名"]
	# page_content: 案情描述  建立向量
	# result：判决结果
	# category 罪名
	for accusation in accusation_list:
		documents.append(Document(
				page_content=text,  # 用于向量化的内容
				metadata={"result": result, "category": accusation}  # 附加信息, 不参与向量计算, 仅作为检索结果的辅助信息
		))
"""
[
    Document(
        page_content="被告人张三于2023年...",  # 案情事实
        metadata={
            "result": "罪名: 故意伤害罪; 刑期: 有期徒刑三年; 罚金: 人民币5000元",
            "category": "故意伤害罪"
        }
    ),
    Document(
        page_content="被告人张三于2023年...",  # 同样的案情事实
        metadata={
            "result": "罪名: 故意伤害罪; 刑期: 有期徒刑三年; 罚金: 人民币5000元",
            "category": "寻衅滋事罪"  # 但分类不同
        }
    ),
    # ... 更多文档
]
"""

# 创建 FAISS 向量数据库
print("开始建索引")
# 传入文档集合和嵌入模型，构建 FAISS 索引
db = FAISS.from_documents(documents, embeddings)
with open("db/law_db", "wb") as f:
	# 将整个 db 对象序列化并保存到文件中
	pickle.dump(db, f)

if __name__ == '__main__':
	with open("db/law_db", "rb") as f:
		# 从文件反序列化并加载数据库对象
		db = pickle.load(f)
	
	# 现在可以进行相似性搜索
	query = "缅北诈骗案件"
	# 进行相似性搜索，返回最相似的 k 个文档，只搜索分类 (category) 为"诈骗"的文档
	similar_docs = db.similarity_search(query, k=10, filter={"category": "诈骗"}, )
	print(similar_docs)
